import json
from fastapi import Body
from typing import List, Optional, Union
from server.api_wrap import api_wrapper, async_api_wrapper, BaseResponse
from control.control import (
    train_model,
    predict,
    select,
    master_controller
)
from loguru import logger
from server.redo_correct import redo_correct_error
import traceback

@async_api_wrapper
async def correct_error_api(
        data: dict = Body(description='错字审查', examples=None)
):
    logger.info(f'错字审查输入: {json.dumps(data, ensure_ascii=False)}')
    try:
        # 先用传统小模型找错别字，速度快。
        result = master_controller.correct_error_text(data['textInfos'])
        # 利用大模型再进行错别字纠正，速度慢
        result = await redo_correct_error(data, result)
    except Exception as e:
        logger.error(e)
        response = {'success': False, 'error': traceback.format_exc(), 'result': None}
        return response
    response = {'success': True, 'error': None, 'result': result}
    logger.info(f'错字审查结果: {json.dumps(response, ensure_ascii=False)}')
    return response


@async_api_wrapper
async def train_model_api(
        data: dict = Body(description='模型训练', examples=None)
):
    logger.info(f'模型训练输入参数: {json.dumps(data, ensure_ascii=False)}')

    result = train_model(data)

    logger.info(f'模型训练结果: {json.dumps(result, ensure_ascii=False)}')
    return result


@async_api_wrapper
async def predict_api(
        data: dict = Body(description='模型训练', examples=None)
):
    logger.info(f'模型训练输入参数: {json.dumps(data, ensure_ascii=False)}')

    result = predict(data)

    logger.info(f'模型训练结果: {json.dumps(result, ensure_ascii=False)}')
    return result


@async_api_wrapper
async def select_api(
        data: dict = Body(description='模型选择', examples=None)
):
    logger.info(f'模型选择输入参数: {json.dumps(data, ensure_ascii=False)}')

    result = select(data)

    logger.info(f'模型选择结果: {json.dumps(result, ensure_ascii=False)}')
    return result


@api_wrapper
def query_similar_text_api(
    queries: List[str] = Body(description='问题的集合', examples=[['你是谁', '吃饭了没']]),
    sorted_texts: List[dict[str, Union[str, int]]] = Body(description='相似文档的查询', examples=[[{'text': '我今天没有吃饭', 'index': 1}, {'text': '我是成龙', 'index': 2}, {'text': '明天多云转阴', 'index': 3}]]),
    task: Optional[str] = Body(default='qa', description='相似文档的查询方式', examples=['qa']),
    top_n: Optional[int] = Body(default=10, description="返回相似的文档数量", examples=[10])
):
    result = master_controller.query_similar_text(queries, sorted_texts, task, top_n)
    return BaseResponse(code=200, status='0', data=result)